# -*- coding: utf8 -*-

# 通过数据库的表结构生成对应的文档。

import sys

from typing import Dict, List, Iterable, Sequence, Tuple, Union, Any
from functools import reduce
from openpyxl import Workbook
from openpyxl.styles.colors import Color
from openpyxl.styles.fills import PatternFill
from openpyxl.utils import datetime
from openpyxl.worksheet.worksheet import Worksheet
from openpyxl.cell.cell import Cell, TYPE_STRING, TYPE_FORMULA, TYPE_NUMERIC, TYPE_BOOL, TYPE_NULL, TYPE_INLINE, TYPE_ERROR, TYPE_FORMULA_CACHE_STRING, WriteOnlyCell
from openpyxl.styles import Font, Alignment
from openpyxl.utils import get_column_letter
from psycopg2 import connect
from datetime import datetime

from .commdef import DataColumn, DataTable
from .utils import find_first, strip_str


def _load_tables_from_db(
    dsn: str,
    verbose: bool
) -> List[DataTable]:
    r''' 从数据库中加载所有的数据库表。

    :param dsn: 数据库连接字符串。
    :param verbose: 是否显示加载过程。

    :return: 已加载的所有数据库表。
    '''
    tables = []
    with connect(dsn=dsn) as conn:
        cur = conn.cursor()
#         cur.execute('''SELECT c1.relname, ARRAY_AGG(c2.relname)
# FROM pg_constraint con
# INNER JOIN pg_class c1 ON c1.oid = con.conrelid
# INNER JOIN pg_class c2 ON c2.oid = con.confrelid
# INNER JOIN pg_namespace n1 ON c1.relnamespace = n1.oid
# INNER JOIN pg_namespace n2 on c2.relnamespace = n2.oid
# WHERE con.contype = 'f'
#   AND c1.relkind = 'r'
#   AND n1.nspname = %s
# GROUP BY c1.relname''', ('public',))
#         table_deps = {}
#         for table_name, table_dep, *_ in cur.fetchall():
#             table_deps[table_name] = set(table_dep)

        cur.execute('''SELECT c.relname, CAST(OBJ_DESCRIPTION(c.relfilenode, 'pg_class') as varchar)
FROM pg_class c
INNER JOIN pg_namespace n ON c.relnamespace = n.oid
WHERE relkind = 'r' AND n.nspname = %s
ORDER BY c.oid''', ('public',))
        all_table_names = []
        for table_name, table_comment in cur:
            all_table_names.append((strip_str(table_name), strip_str(table_comment)))

        # 按照依赖关系计算次序。

        if verbose:
            print('Found {:d} tables'.format(len(all_table_names)))

        for table_name, table_comment in all_table_names:
            if verbose:
                print('Process {!r:} ...'.format(table_name))

            cols = []

            # 查询列定义。
            cur.execute('''SELECT a.attname,
    t.typname,
    a.atttypmod,
    col_description(a.attrelid, a.attnum),
    a.attnotnull,
    pg_get_expr(d.adbin, d.adrelid, false),
    con.contype IS NOT NULL
FROM pg_attribute a
INNER JOIN pg_class c ON a.attrelid = c.oid
INNER JOIN pg_type t ON a.atttypid = t.oid
LEFT JOIN pg_attrdef d ON d.adrelid = c.oid AND d.adnum = a.attnum
LEFT JOIN pg_constraint con ON con.conrelid = c.oid AND con.conkey[1] = a.attnum AND con.contype = 'p'
WHERE c.relname = %s
AND a.attnum > 0
ORDER BY a.attnum''', (table_name,))
            col_results = cur.fetchall()

            for col_name, col_type, col_type_len, col_comment, col_nullable, col_default, col_primary_key, *_ in col_results:
                # 规范化类型名。
                if col_type == 'int8':
                    col_type = 'bigint'
                elif col_type == 'int4':
                    col_type = 'int'
                if col_default and col_default.startswith('nextval('):
                    col_default = ''
                    if col_type == 'integer':
                        col_type = 'serial'
                    elif col_type == 'bigint':
                        col_type = 'bigserial'

                # 计算宽度。
                if col_type_len > 0x0ffff:
                    col_size = (col_type_len & 0x0ff0000) >> 16
                    col_scale = (col_type_len & 0x0ffff) - 4
                elif col_type_len > 0:
                    col_size = col_type_len - 4
                    col_scale = None
                else:
                    col_size = col_scale = None

                # 如果存在主键约束，那么该列非空。
                if col_primary_key:
                    col_nullable = False

                cols.append(DataColumn(col_name, col_type, col_size, col_scale, col_comment,
                                       col_nullable, col_default, '', '', col_primary_key, '', ''))

            # 查询约束条件（只包含检查约束和外键约束）。
            cur.execute('''SELECT a.attname, con.contype, pg_get_expr(con.conbin, c.oid, false), fc.relname
FROM pg_constraint con
INNER JOIN pg_attribute a ON a.attnum = con.conkey[1]
INNER JOIN pg_class c ON a.attrelid = c.oid AND con.conrelid = c.oid
LEFT JOIN pg_class fc ON con.confrelid = fc.oid
LEFT JOIN pg_attribute a2 ON a2.attrelid = c.oid AND a2.attnum = ANY(con.conkey)
WHERE c.relname = %s
AND a.attnum > 0
AND con.contype IN ('c', 'f')
ORDER BY a.attnum''', (table_name,))
            con_result = cur.fetchall()

            for col_name, contype, con_value, fk_table, *_ in con_result:
                col = find_first(lambda _, col_: col_.col_name == col_name, cols)
                if col is None:
                    continue
                if contype == 'u':
                    # 发现唯一约束。
                    col.col_unique = ','.join(con_value)
                elif contype == 'c':
                    # 发现检查约束。
                    # 从约束表达式中去掉括号和当前列名。
                    trans = str.maketrans({'(': '', ')': ''})
                    col.col_check = ','.join(con_item.translate(trans).replace(col.col_name, '').strip() for con_item in con_value.split('AND') if con_item)
                elif contype == 'f':
                    # 发现外键约束。
                    col.col_foreign_key = fk_table

            # 检查唯一约束。
            cur.execute('''SELECT ARRAY_AGG(a.attname)
FROM pg_constraint con
INNER JOIN pg_attribute a ON a.attnum = ANY(con.conkey)
INNER JOIN pg_class c ON a.attrelid = c.oid AND con.conrelid = c.oid
WHERE c.relname = %s
AND a.attnum > 0
AND con.contype = 'u'
GROUP BY con.oid
ORDER BY con.oid''', (table_name,))
            unique_result = cur.fetchall()

            for rn, row in enumerate(unique_result):
                uniques = row[0]
                for col_name in uniques:
                    col = find_first(lambda _, col_: col_.col_name == col_name, cols)
                    if col is None:
                        continue
                    col.col_unique = str(rn + 1)

            # 检查索引。
            cur.execute('''SELECT ARRAY_AGG(a.attname)
FROM pg_index ind
INNER JOIN pg_attribute a ON a.attnum = ANY(ind.indkey)
INNER JOIN pg_class c ON a.attrelid = c.oid AND ind.indrelid = c.oid
WHERE c.relname = %s
AND NOT ind.indisunique
GROUP BY ind.indexrelid
ORDER BY ind.indexrelid''', (table_name,))
            idx_result = cur.fetchall()

            for rn, row in enumerate(idx_result):
                idx = row[0]
                for col_name in idx:
                    col = find_first(lambda _, col_: col_.col_name == col_name, cols)
                    if col is None:
                        continue
                    col.col_index = str(rn + 1)

            tables.append(DataTable(table_name, table_comment, cols))

    # 表的依赖关系。
    return tables


def _create_header(
    sheet: Worksheet,
    *args: str
) -> None:
    r''' 创建表格的表头。

    :param *args: 表头的标题。
    '''
    tf = Font(bold=True)
    tpf = PatternFill(patternType='solid', fgColor=Color(rgb='a0a0a0'))
    for ai, arg in enumerate(args):
        cell: Cell = sheet.cell(1, ai + 1, value=arg)
        cell.alignment = Alignment(horizontal='center', vertical='center', wrap_text=False)
        cell.font = tf
        cell.fill = tpf


def _adjust_col_size(
    sheet: Worksheet
) -> None:
    r''' 自动调整列宽。

    :param sheet: 需要调整列宽的工作簿。
    '''
    def _text_size(s) -> int:
        return reduce(lambda l, ch: l + 1 if ch.isascii() else l + 2, s, 2)

    col_d = sheet.column_dimensions
    dims = {}
    for row in sheet.rows:
        for cell in row:
            v = cell.value
            cw = _text_size(str(v)) if v is not None else 0
            if cell.value:
                dims[cell.column] = max(dims.get(cell.column, 0), cw)
    for col, width in dims.items():
        col_d[get_column_letter(col)].width = width


def _dump_doc(
    tables: Sequence[DataTable],
    file=sys.stdout
) -> None:
    wb = Workbook()
    wb.properties.title = '数据字典'
    wb.properties.creator = 'Nau'
    wb.properties.created = datetime.now()
    if wb.active:
        wb.remove_sheet(wb.active)

    # 输出目录页。
    index_sheet = wb.create_sheet(title='目录')
    _create_header(index_sheet, '名称', '说明')
    for table_i, table in enumerate(tables):
        hcell: Cell = index_sheet.cell(table_i + 2, 1, value=table.name)
        hcell.hyperlink = '#' + table.name + '!A1'
        hcell.style = 'Hyperlink'
        index_sheet.cell(table_i + 2, 2, value=table.comment)

    _adjust_col_size(index_sheet)

    # 每个数据库表对应一个Sheet。
    for table in tables:
        sheet = wb.create_sheet(title=table.name)
        _create_header(sheet, '列名', '类型', '长度', '小数', '说明', '可空', '默认值', '唯一',	'检查',	'主键', '外键',	'索引')
        rn = 2
        for col in table.cols:
            sheet.cell(rn, 1, value=col.col_name)
            sheet.cell(rn, 2, value=col.col_type)
            sheet.cell(rn, 3, value=col.col_size)
            sheet.cell(rn, 4, value=col.col_scale)
            sheet.cell(rn, 5, value=col.col_comment)
            sheet.cell(rn, 6, value='是' if col.col_nullable else '')
            sheet.cell(rn, 7, value=col.col_default)
            sheet.cell(rn, 8, value=col.col_unique)
            sheet.cell(rn, 9, value=col.col_check)
            sheet.cell(rn, 10, value='是' if col.col_primary_key else '')
            sheet.cell(rn, 11, value=col.col_foreign_key)
            sheet.cell(rn, 12, value=col.col_index)
            rn += 1

        _adjust_col_size(sheet)

    wb.save(file)


def execute(
    dsn: str,
    doc_path: str,
    verbose: bool = False
) -> None:
    r''' 数据库连接字符串。
    '''
    tables = _load_tables_from_db(dsn, verbose)

    with open(doc_path, 'wb') as f:
        _dump_doc(tables, file=f)
